/*
 * =============================================================
 * Simple MEX wrapper for maxflow algorithm.
 *
 * Should be called through MinCut() (defined in MinCut.m)
 *
 * @note This is used to segment an image. It is a modified version
 * of the mex file by micham23 (January 2006) that was used for panorama stitching.
 * @note - Used the adjacency_list data structure, with MODIFIED 
 * 		 graph.h ('captype' and 'flowtype' changed to 'double')
 *
 * Author: yairmov@cs.huji.ac.il
 * =============================================================
 */

#include "mex.h"
#include "graph.h"


/**
 * Thin wrapper for the maxflow algorithm.
 */
void stich(
	int M, int N, 
	double* horizontalEdges, double* verticalEdges, 
	double* sourceWeights, double* sinkWeights,
	bool* result) {

	// Init.
	Graph* g = new Graph();
	Graph::node_id* nodes = new Graph::node_id[M*N];

	// Predeclare for-loop vars
	// @note Yes, this is ugly, but done for cross-compiler support
	int m, n;
		
	// Process horizontal and vertical edges...
	for (m = 0; m < M; ++m) {
		for (n = 0; n < N; ++n) {
			// Pre-calculate absolute index
			int index = n*M+m;

			// Add node
			nodes[index] = g->add_node();

			// Add horizontal edge
			if (n > 0) {
				int hIndex = index;

				// Use weight from the neighbor to our left
				g->add_edge(
					nodes[index], nodes[index-M], 
					horizontalEdges[hIndex-M], horizontalEdges[hIndex-M]);
			}
				
			// Add vertical edge
			if (m > 0) {
				// verticalEdges matrix is missing the last row
				int vIndex = n*(M-1)+m;

				// Use weight from the neighbor above us
				g->add_edge(
					nodes[index], nodes[index-1], 
					verticalEdges[vIndex-1], verticalEdges[vIndex-1]);
			}
			// set weights to source node and sink node.
			g->set_tweights(nodes[index], sourceWeights[index], sinkWeights[index]);
		}
	}
									

	// Just do it!
	Graph::flowtype flow = g->maxflow();
	//mexPrintf("findMinCut: Max flow is %.6lf\n", flow);

	// Create result segmentation
	for (m = 0; m < M; ++m) {
		for (n = 0; n < N; ++n) {
			// Pre-calculate absolute index
			int index = n*M+m;

			result[index] = (g->what_segment(nodes[index]) == Graph::SOURCE);
		}
	} 

	// Clean-up
	delete[] nodes;
	delete g;
}

// Constants for parameter indices
#define IS_FRIEND_PARAM_INDEX			0
#define HORIZONTAL_EDGES_PARAM_INDEX		1
#define VERTICAL_EDGES_PARAM_INDEX		2
#define SOURCE_NODES_PARAM_INDEX		3
#define	SINK_NODES_PARAM_INDEX			4

// Constants for return value indices
#define SEGMENTATION_RETVAL_INDEX		0

/** 
 * Standard MEX entry-point.
 * 
 * Just a thin layer to convert from/to matlab types/interface.
 */
void mexFunction(
	int nlhs, mxArray *plhs[], 
	int nrhs, const mxArray *prhs[]) {

	// Make sure that our caller is the matlab minCut() function.
	// We check the number of input and output parameters, and that the 
	// 'isFriend' first parameter is true. This is done just to avoid people 
	// calling us by mistake.
	// @note We rely on minCut()'s syntax checking, and do NOT perform any
	// 		 more checks
	if ((nrhs != 5) || (nlhs > 1) 
					|| (! mxIsLogicalScalarTrue(prhs[IS_FRIEND_PARAM_INDEX]))) {
		mexErrMsgTxt("mexFindMinCut: Do NOT use directly. "
						"Use MinCut() instead!");
	}

	// @note M is taken from first edges argument, 
	// 	     while N is taken from the second.
	int M = mxGetM(prhs[HORIZONTAL_EDGES_PARAM_INDEX]);
	int N = mxGetN(prhs[VERTICAL_EDGES_PARAM_INDEX]);

	// Create matrix for the return argument.
	plhs[SEGMENTATION_RETVAL_INDEX] = mxCreateLogicalMatrix(M, N);

	// Just do it!
	stich(M, N, 
		  mxGetPr(prhs[HORIZONTAL_EDGES_PARAM_INDEX]), 
		  mxGetPr(prhs[VERTICAL_EDGES_PARAM_INDEX]), 
		  mxGetPr(prhs[SOURCE_NODES_PARAM_INDEX]), 
		  mxGetPr(prhs[SINK_NODES_PARAM_INDEX]),
		  mxGetLogicals(plhs[SEGMENTATION_RETVAL_INDEX]));
}

